package com.central.scheduler.master.runner.task;

import com.central.common.constant.CommonConstant;
import com.central.common.entity.*;
import com.central.common.enums.DbType;
import com.central.common.enums.PrivilegeType;
import com.central.common.enums.TaskType;
import com.central.common.model.*;
import com.central.scheduler.common.Constants;
import com.central.scheduler.common.utils.HadoopUtils;
import com.central.scheduler.common.utils.JSONUtils;
import com.central.scheduler.service.bean.SpringApplicationContext;
import com.central.scheduler.service.process.ProcessService;
import com.central.scheduler.spi.task.request.DataxTaskExecutionContext;
import com.central.scheduler.service.queue.entity.TaskExecutionContext;
import com.central.scheduler.spi.task.request.FlinkSqlTaskExecutionContext;
import org.apache.commons.collections4.CollectionUtils;
import com.central.scheduler.dao.model.ProcessInstance;
import com.central.scheduler.dao.model.TaskInstance;
import com.central.scheduler.master.builder.TaskExecutionContextBuilder;
import com.central.scheduler.spi.task.request.SQLTaskExecutionContext;
import com.central.scheduler.spi.task.request.UdfFuncRequest;
import com.central.scheduler.spi.utils.PropertyUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;

/**
 * @author Tindy
 * @date 2021/12/18
 * @describe
 */
public abstract class BaseTaskProcessor implements ITaskProcessor {

    protected Logger logger = LoggerFactory.getLogger(getClass());

    protected boolean killed = false;

    protected boolean paused = false;

    protected boolean timeout = false;

    protected TaskInstance taskInstance = null;

    protected ProcessInstance processInstance;

    protected ProcessService processService = SpringApplicationContext.getBean(ProcessService.class);

    /**
     * pause task, common tasks donot need this.
     *
     * @return
     */
    protected abstract boolean pauseTask();

    /**
     * kill task, all tasks need to realize this function
     *
     * @return
     */
    protected abstract boolean killTask();

    /**
     * task timeout process
     * @return
     */
    protected abstract boolean taskTimeout();

    @Override
    public void run() {
    }

    @Override
    public boolean action(TaskAction taskAction) {

        switch (taskAction) {
            case STOP:
                return stop();
            case PAUSE:
                return pause();
            case TIMEOUT:
                return timeout();
            default:
                logger.error("unknown task action: {}", taskAction.toString());

        }
        return false;
    }

    protected boolean timeout() {
        if (timeout) {
            return true;
        }
        timeout = taskTimeout();
        return timeout;
    }

    /**
     * @return
     */
    protected boolean pause() {
        if (paused) {
            return true;
        }
        paused = pauseTask();
        return paused;
    }

    protected boolean stop() {
        if (killed) {
            return true;
        }
        killed = killTask();
        return killed;
    }

    @Override
    public String getType() {
        return null;
    }

    /**
     * get TaskExecutionContext
     *
     * @param taskInstance taskInstance
     * @return TaskExecutionContext
     */
    protected TaskExecutionContext getTaskExecutionContext(TaskInstance taskInstance) {
        processService.setTaskInstanceDetail(taskInstance);

        taskInstance.setResources(getResources(taskInstance));

        SQLTaskExecutionContext sqlTaskExecutionContext = new SQLTaskExecutionContext();
        DataxTaskExecutionContext dataxTaskExecutionContext = new DataxTaskExecutionContext();
        FlinkSqlTaskExecutionContext flinkSqlTaskExecutionContext = new FlinkSqlTaskExecutionContext();


        // SQL task
        if (TaskType.SQL.equals(taskInstance.getTaskType())) {
            setSQLTaskRelation(sqlTaskExecutionContext, taskInstance);
        }

        // DATAX task
        if (TaskType.DATAX.equals(taskInstance.getTaskType())) {
            setDataxTaskRelation(dataxTaskExecutionContext, taskInstance);
        }

        // FLINK_SQL task
        if (TaskType.FLINK_SQL.equals(taskInstance.getTaskType())) {
            setFlinkSqlTaskRelation(flinkSqlTaskExecutionContext, taskInstance);
        }

        return TaskExecutionContextBuilder.get()
                .buildTaskInstanceRelatedInfo(taskInstance)
                .buildProcessInstanceRelatedInfo(taskInstance.getProcessInstance())
                .buildProcessDefinitionRelatedInfo(taskInstance.getProcessDefine())
                .buildSQLTaskRelatedInfo(sqlTaskExecutionContext)
                .buildDataxTaskRelatedInfo(dataxTaskExecutionContext)
                .buildFlinkSqlTaskRelatedInfo(flinkSqlTaskExecutionContext)
                .create();
    }

    /**
     * set flink sql task relation
     *
     * @param flinkSqlTaskExecutionContext flinkSqlTaskExecutionContext
     * @param taskInstance taskInstance
     */
    private void setFlinkSqlTaskRelation(FlinkSqlTaskExecutionContext flinkSqlTaskExecutionContext, TaskInstance taskInstance) {
        FlinkSqlTaskContent taskContent = JSONUtils.parseObject(taskInstance.getTaskContent(), FlinkSqlTaskContent.class);
        boolean dataPrivilegeEnable = PropertyUtils.getBoolean(Constants.DATA_PRIVILEGE_ENABLE, false);

        Map<String, FlinkTable> tableMap = taskContent.getTableMap();
        Map<String, Set<PrivilegeType>> allTablePrivilege =new HashMap<>();
        for (FlinkTable table : tableMap.values()) {
            Long datasourceId = table.getDatasourceId();
            DatasourceInfo datasource = processService.findDataSourceById(datasourceId);
            table.setDatasource(datasource);
            TableInfo tableInfo = processService.getTableInfo(table.getDatasourceId(), table.getDbName(), table.getTableName());
            table.setTableInfo(tableInfo);
            if(dataPrivilegeEnable){
                Set<String> tables = new HashSet<>();
                tables.add(table.getDbName()+"."+table.getTableName());
                Map<String, Set<PrivilegeType>> tablePrivilege = processService.getAllTablePrivilege(datasourceId, taskInstance.getProId(), taskInstance.getUserId(),tables);
                allTablePrivilege.putAll(tablePrivilege);
            }
        }
        flinkSqlTaskExecutionContext.setTablePrivileges(allTablePrivilege);
        if(dataPrivilegeEnable){
            List<Map<String, Set<PrivilegeType>>> needPrivileges = taskInstance.getTaskDefine().getDataRelateds().stream().map(x -> x.getNeedPriveleges()).collect(Collectors.toList());
            flinkSqlTaskExecutionContext.setNeedPrivileges(needPrivileges);
        }
        flinkSqlTaskExecutionContext.setTableMap(tableMap);


        // whether udf type
        boolean udfTypeFlag = !CollectionUtils.isEmpty(taskContent.getUdfIds());
        ArrayList<UdfFuncRequest> udfRequests = new ArrayList<>();
        if (udfTypeFlag) {
            List<Long> udfIds = taskContent.getUdfIds();

            List<UdfDefined> udfs = processService.queryUdfsByIds(udfIds);
            for (UdfDefined udf : udfs) {
                UdfFuncRequest udfFuncRequest = new UdfFuncRequest();
                udfFuncRequest.setFuncName(udf.getFuncName());
                udfFuncRequest.setClassName(udf.getClassName());
                udfFuncRequest.setResourceUrl(udf.getResourceUrl());
                udfFuncRequest.setResourceName(udf.getResourceName());
                udfRequests.add(udfFuncRequest);
            }
        }
        flinkSqlTaskExecutionContext.setUdfs(udfRequests);


        //set queue
        if(taskContent.getQueueId()!=null){
            QueueInfo queue = processService.getQueue(taskContent.getQueueId());
            if(queue!=null){
                flinkSqlTaskExecutionContext.setQueue(queue.getName());
            }
        }
        //set user keytab url
        flinkSqlTaskExecutionContext.setUserKtUrl(processService.getUserKtUrl(taskInstance.getProName(),taskInstance.getUsername()));
    }


    /**
     * set datax task relation
     *
     * @param dataxTaskExecutionContext dataxTaskExecutionContext
     * @param taskInstance taskInstance
     */
    protected void setDataxTaskRelation(DataxTaskExecutionContext dataxTaskExecutionContext, TaskInstance taskInstance) {
        DataxTaskContent taskContent = JSONUtils.parseObject(taskInstance.getTaskContent(),DataxTaskContent.class);
        Long readerDatasourceId = taskContent.getReader().getDatasourceId();
        Long writerDatasourceId = taskContent.getWriter().getDatasourceId();
        DatasourceInfo readerDatasource = processService.findDataSourceById(readerDatasourceId);
        DatasourceInfo writerDatasource = processService.findDataSourceById(writerDatasourceId);
        boolean dataPrivilegeEnable = PropertyUtils.getBoolean(Constants.DATA_PRIVILEGE_ENABLE, false);
        if(dataPrivilegeEnable){
            DataRelated dataRelated = taskInstance.getTaskDefine().getDataRelateds().get(0);
            Set<String> readerTables = new HashSet<>();
            readerTables.addAll(dataRelated.getSourceTables());
            Set<String> writerTables = new HashSet<>();
            writerTables.add(dataRelated.getTargetTable());
            Map<String, Set<PrivilegeType>> readerTablePrivileges = processService.getAllTablePrivilege(readerDatasourceId, taskInstance.getProId(), taskInstance.getUserId(),readerTables);
            dataxTaskExecutionContext.setReaderTablePrivileges(readerTablePrivileges);
            Map<String, Set<PrivilegeType>> writerTablePrivileges = processService.getAllTablePrivilege(writerDatasourceId, taskInstance.getProId(), taskInstance.getUserId(),writerTables);
            dataxTaskExecutionContext.setWriterTablePrivileges(writerTablePrivileges);
        }
        dataxTaskExecutionContext.setReaderDatasource(readerDatasource);
        dataxTaskExecutionContext.setWriterDatasource(writerDatasource);
    }

    /**
     * set SQL task relation
     *
     * @param sqlTaskExecutionContext sqlTaskExecutionContext
     * @param taskInstance taskInstance
     */
    private void setSQLTaskRelation(SQLTaskExecutionContext sqlTaskExecutionContext, TaskInstance taskInstance) {
        SqlTaskContent sqlTaskContent = JSONUtils.parseObject(taskInstance.getTaskContent(),SqlTaskContent.class);
        long datasourceId = sqlTaskContent.getDatasourceId();
        DatasourceInfo datasource = processService.findDataSourceById(datasourceId);
        boolean dataPrivilegeEnable = PropertyUtils.getBoolean(Constants.DATA_PRIVILEGE_ENABLE, false);
        if(dataPrivilegeEnable && datasource!=null){
            if(datasource.getIsSys() && datasource.getDatasourceType()== DbType.HIVE){
                datasource.setUsername(CommonConstant.KA_PREFIX+taskInstance.getProName()+"_"+taskInstance.getUsername()+"/"+CommonConstant.KA_PREFIX+taskInstance.getProName());
                datasource.setKeytabUrl(processService.getUserKtUrl(taskInstance.getProName(),taskInstance.getUsername()));
                datasource.setDbName(sqlTaskContent.getDbName());
            }else{
                List<DataRelated> dataRelateds = taskInstance.getTaskDefine().getDataRelateds();
                HashSet<String> tables = new HashSet<>();
                for (DataRelated dataRelated : dataRelateds) {
                    tables.addAll(dataRelated.getSourceTables());
                    if (dataRelated.getTargetTable() != null) {
                        tables.add(dataRelated.getTargetTable());
                    }
                }
                Map<String, Set<PrivilegeType>> allTablePrivilege = processService.getAllTablePrivilege(datasourceId, taskInstance.getProId(), taskInstance.getUserId(),tables);
                sqlTaskExecutionContext.setTablePrivileges(allTablePrivilege);
                List<Map<String, Set<PrivilegeType>>> needPrivileges = dataRelateds.stream().map(x -> x.getNeedPriveleges()).collect(Collectors.toList());
                sqlTaskExecutionContext.setNeedPrivileges(needPrivileges);
            }
        }

        sqlTaskExecutionContext.setDatasource(datasource);

        sqlTaskExecutionContext.setDefaultFS(HadoopUtils.getInstance().getDefaultFS());

        // whether udf type
        boolean udfTypeFlag = !CollectionUtils.isEmpty(sqlTaskContent.getUdfIds());

        if (udfTypeFlag) {
            List<Long> udfIds = sqlTaskContent.getUdfIds();

            List<UdfDefined> udfs = processService.queryUdfsByIds(udfIds);
            ArrayList<UdfFuncRequest> udfRequests = new ArrayList<>();
            for (UdfDefined udf : udfs) {
                UdfFuncRequest udfFuncRequest = new UdfFuncRequest();
                udfFuncRequest.setFuncName(udf.getFuncName());
                udfFuncRequest.setClassName(udf.getClassName());
                udfFuncRequest.setResourceUrl(udf.getResourceUrl());
                udfFuncRequest.setResourceName(udf.getResourceName());
                udfRequests.add(udfFuncRequest);
            }
            sqlTaskExecutionContext.setUdfs(udfRequests);
        }

        //set engine
        if(sqlTaskContent.getEngineType()!=null ){
            sqlTaskExecutionContext.setEngine(sqlTaskContent.getEngineType());
        }
        //set queue
        if(sqlTaskContent.getQueueId()!=null){
            QueueInfo queue = processService.getQueue(sqlTaskContent.getQueueId());
            if(queue!=null){
                sqlTaskExecutionContext.setQueue(queue.getName());
            }
        }
        // set setParam
        sqlTaskExecutionContext.setSetParams(sqlTaskContent.getSetParams());
    }


    protected List<ResourceInfo> getResources(TaskInstance taskInstance) {
        if(taskInstance.getResourceIds()!=null){
            List<ResourceInfo> resources = processService.getResources(taskInstance.getResourceIds());
            return resources;
        }else {
            return new ArrayList<>();
        }
    }
}
